import cv2
import numpy as np

# 图像分割 并 Mnist 化
class ImageDeal:
    # 传入图片
    def __init__(self, imagePath) -> None:
        self.image = cv2.imread(imagePath, cv2.IMREAD_GRAYSCALE)

        self._accessBinary()

    # 反向灰度图
    def _accessPiexl(self):
        height = self.image.shape[0]
        width = self.image.shape[1]

        for i in range(height):
            for j in range(width):
                self.image[i][j] = 255 - self.image[i][j]

    # 反向二值化图像
    def _accessBinary(self, threshold=128):
        # 颠倒
        self._accessPiexl()
        # 二值化
        _, self.image = cv2.threshold(self.image, threshold, 0, cv2.THRESH_TOZERO)

    # 寻找边缘，返回边框的左上与右下，并分割图像
    def getImageSagList(self, minArea=30):
        # 边缘 最外部 不近似
        contours, _ = cv2.findContours(self.image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
        # 存储边框顶点
        borders = []

        for contour in contours:
            # 将边缘拟合为边框
            x, y, w, h = cv2.boundingRect(contour)
            # 边框足够大
            if w * h > minArea:
                border = [(x, y), (x + w, y + h)]
                borders.append(border)

        # 根据图片的 x 轴排序
        borders.sort(key= lambda x: ((x[0][0] + x[1][0]) + (x[0][1] + x[1][1]) * 5))

        # 存储分割的图像数据
        imageSagList = []
        for border in borders:
            borderImage = self.image[border[0][1] : border[1][1], border[0][0] : border[1][0]]
            imageSagList.append(borderImage)

        return imageSagList

    # 根据顶点将图像分割，并将图片格式转化为 MNIST 格式
    @staticmethod
    def transMNIST(imageSagList, size=(28, 28)):
        # 无符号图像数组
        targetImageList = np.zeros((len(imageSagList), size[0], size[1], 1), dtype="uint8")
        # 遍历切割图片
        for i, imageSag in enumerate(imageSagList):
            # 根据最大边缘拓展像素
            extendPiexl = max(imageSag.shape) // 2
            targetImage = cv2.copyMakeBorder(
                imageSag, extendPiexl, extendPiexl, extendPiexl, extendPiexl, cv2.BORDER_CONSTANT
            )
            # 图像缩放
            targetImage = cv2.resize(targetImage, size)
            # 拓展维度
            targetImage = np.expand_dims(targetImage, axis=-1)
            targetImageList[i] = targetImage

        return targetImageList

# 测试
if __name__ == "__main__":
    imagePath = "./image/test.png"

    ID = ImageDeal(imagePath)
    imageSagList = ID.getImageSagList()

    targetImageList = ID.transMNIST(imageSagList)

    cv2.imshow("1.png", targetImageList[0])
    cv2.waitKey()
    cv2.destroyAllWindows()
